package com.han.agent.controller;

import com.han.agent.advisor.LoggerAdvisor;
import com.han.agent.entity.Movie;
import com.han.agent.tools.DateTool;
import io.micrometer.common.util.StringUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * @author han
 */
@RestController
@RequestMapping("/chat")
public class ChatController {

    private final OpenAiChatModel chatModel;
    private final ChatClient chatClient;
    private final ChatClient nameChatClient;
    private final ChatClient writerChatClient;
    private final ChatClient memoryChatClient;

    /**
     * 同一个chatClientBuilder构建多个ChatClient，defaultSystem会覆盖
     */
    public ChatController(ChatClient.Builder chatClientBuilder, OpenAiChatModel chatModel) {
        this.chatModel = chatModel;
        this.chatClient = chatClientBuilder.build();
        this.nameChatClient = ChatClient.builder(chatModel)
                .defaultSystem("你是一个老道士，起了很多富有中国文化内涵的男孩名字，给我一个你觉得最好的就行了")
                .build();

        this.writerChatClient = ChatClient.builder(chatModel)
                .defaultSystem("你是一个非常牛逼且非常会写装逼小说的小说家，根据输入的文章题目写一篇200字以内的短篇小说")
                .build();

        this.memoryChatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(MessageWindowChatMemory.builder().build()).build())
                .build();
    }

    @GetMapping("/ai/generate")
    public String generate(@RequestParam(required = false, defaultValue = "告诉我一个笑话") String message) {
        return this.chatModel.call(message);
    }

    @GetMapping("/naming/generate")
    public String naming(@RequestParam(required = false, defaultValue = "起个姓韩的名字") String userInput) {
        return this.nameChatClient.prompt()
                .user(userInput)
                .call()
                .content();
    }

    @GetMapping("/writer/generate")
    public String writer(@RequestParam(required = false, defaultValue = "仿照斗破苍穹来一篇") String userInput) {
        return this.writerChatClient.prompt()
                .user(userInput)
                .call()
                .content();
    }

    /**
     * 流式响应Flux
     */
    @GetMapping(value = "/writer/generate/flux", produces = "text/html;charset=UTF-8")
    public Flux<String> streamChat(@RequestParam(required = false, defaultValue = "仿照斗破苍穹来一篇") String userInput) {
        return writerChatClient.prompt()
                .user(userInput)
                .stream()
                .content();
    }

    /**
     * 流式响应SSE
     */
    @GetMapping("/writer/generate/sse")
    public SseEmitter sseChat(@RequestParam(required = false, defaultValue = "仿照斗破苍穹来一篇") String userInput) {
        SseEmitter sseEmitter = new SseEmitter() {
            @Override
            protected void extendResponse(ServerHttpResponse outputMessage) {
                HttpHeaders headers = outputMessage.getHeaders();
                headers.setContentType(new MediaType("text", "event-stream", StandardCharsets.UTF_8));
            }
        };

        Flux<String> stream = writerChatClient
                .prompt()
                .user(userInput)
                .stream()
                .content();

        stream.subscribe(token -> {
            try {
                sseEmitter.send(token);
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }
        }, sseEmitter::completeWithError, sseEmitter::complete);
        return sseEmitter;
    }

    /**
     * 结构化输出
     */
    @GetMapping("/ai/entity")
    public Movie entity() {
        return this.chatClient.prompt()
                .user("给我周星驰导演的电影信息")
                .call()
                .entity(Movie.class);
    }

    /**
     * 结构化输出 - 列表
     */
    @GetMapping("/ai/entities")
    public List<Movie> entities() {
        BeanOutputConverter<List<Movie>> outputConverter = new BeanOutputConverter<>(
                new ParameterizedTypeReference<>() {
                });

        String format = outputConverter.getFormat();
        String template = """
                给我周星驰的三部电影信息
                {format}
                """;

        Prompt prompt = PromptTemplate.builder().template(template).variables(Map.of("format", format)).build().create();
        String content = chatClient.prompt(prompt).call().content();

        if (StringUtils.isBlank(content)) {
            return Collections.emptyList();
        }

        return outputConverter.convert(content);
    }

    @GetMapping("/ai/advisor")
    public String test(@RequestParam(required = false, defaultValue = "请给我写一首七言绝句") String userInput) {
        ChatClient.Builder builder = ChatClient.builder(chatModel);
        ChatClient chatClient = builder
                .defaultAdvisors(new LoggerAdvisor())
                .build();

        return chatClient.prompt()
                .user(userInput)
                .call()
                .content();
    }

    /**
     * 记忆对话 - 同一个chatId即可
     */
    @GetMapping("/ai/memory")
    public String chatMemoryChat(@RequestParam("chatId") String chatId, @RequestParam("message") String message) {
        return memoryChatClient
                .prompt()
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .user(message)
                .call()
                .content();
    }

    /**
     * 日期工具调用
     */
    @GetMapping("/ai/dateTool")
    public String tool(@RequestParam(required = false, defaultValue = "现在的时间是什么") String message) {
        return chatClient
                .prompt()
                .user(message)
                .tools(new DateTool())
                .call()
                .content();
    }

    /**
     * 本地资源工具调用
     */
    @GetMapping("/ai/localResource")
    public String localResource(@RequestParam(required = false, defaultValue = "test文档.txt的内容是什么") String message) {
        return chatClient
                .prompt()
                .user(message)
                .tools(new DateTool())
                .call()
                .content();
    }

}
